/**
 \file DssFragmentation.cxx
 Implementation of class DssFragmentation.
 
 \author Thomas Burton 
 \date 11/18/11
 */

#include <cmath>
#include <list>

#include <boost/foreach.hpp>
#include <boost/multi_array.hpp>

#include <TMath.h>
#include <TParticlePDG.h>
#include <TRandom3.h>
#include <TROOT.h>
#include <TVector3.h>
#include <TVirtualPad.h>

#include "tppmc/DssFragmentation.h"
#include "tppmc/KkpParton.h"

extern "C" {
   
   // Access to FORTRAN fDSS subroutine (see file fDSS.f).
   void fdss_(int* hadron, int* charge, int* order, double* z,
              double* scale, int* parton, double* result);
   
   // Access to FORTRAN FRAGINI COMMON block (see file fDSS.f).
   extern struct {
      double fini; 
   } fragini_;
   
} // extern "C"

namespace {

enum {
   DUMMY = 0,
   SCALE = 1,
   PARTON = 2,
   ORDER = 3,
   HADRON = 4,
   CHARGE = 5
};

/*
 Wrapper around Fortran DSS subroutine (see fDSS.f).
 Compatible with ROOT TF1.
 x[0] is the z value at which to evaluate the function.
 Depends on six parameters, par[0] to par[5]:
 0: Fragmentation function type: DSS or DSS + Collins
 1: Q at which to evaluate (not Q^2)
 2: Parton type, following KKP convention (see fDSS.f)
 3: LO or NLO function
 4: Hadron species (pion, kaon, proton or charged hadrons)
 5: Hadron charge (+, - or average of +/-)
 */
Double_t dssFragRootWrapper(Double_t* x, Double_t* par) {
   // Note implicit type conversions from par[...] to integers
   // for some parameters.
   double z = x[0];
   double scale = pow(par[SCALE], 2.); // Parameter stores Q, want Q^2
   int parton = par[PARTON];
   int order = par[ORDER]; // LO or NLO
   int hadron = par[HADRON]; // species
   int charge = par[CHARGE]; // charge
   double dss; // Value of the fragmentation function
   
   // Set the "first" flag in the DSS routine.
   static bool first = true;
   
   if(first) {
      fragini_.fini = 0;
      first = false; 
   } // if
   
   // If below charm/bottom limits then fragment as u/ubar
   if(((CQUARK == parton) or (CBAR == parton)) and (scale < (1.65 * 1.65))) {
      if(CQUARK == parton)
         parton = UQUARK;
      else
         parton = UBAR;
   } // if
   
   if(((BQUARK == parton) or (BBAR == parton)) and (scale < (4.85 * 4.85))) {
      /**
       \remark Shouldn't these assignments be the other way round
       for correct charge?
       */
      if(BQUARK == parton)
         parton = UQUARK;
      else
         parton = UBAR;
   } // if
   
   // Call the FORTRAN routine
   fdss_(&hadron,&charge,&order,&z,&scale,&parton,&dss);
   
   // The NLO set for protons (hadron==3) from gluons (parton==0)
   // returns small negative values for low q^2 (~1) and small z (~0.05).
   // Deal with this by setting the ff to zero. 
   
   if(dss < 0.) {
      if(not ((3 == hadron) and (0 == parton) and (z < 0.06))) {
         std::cerr << "DSS WARNING, D_h<0.0 (" << dss << ") for z = "
         << x[0] << " qs = " << par[1] << " parton = " << par[2] << " hadron = " 
         << par[4] << " order = " << par[3] << " charge = " << par[5] 
         << std::endl;
      } // if
      else {
         dss = 0.0; 
      } // else
   } // if
   
   return dss;
}

} // namespace

namespace tppmc {

TF1* DssFragmentation::Function() {
   if(not mFunction) {
      int count(0);
      TString namebase("tppmc_dss_frag_function");
      TString name;
      do {
         name = namebase;
         name += count;
         ++count;
      } // do
      while(gROOT->FindObject(name));
      
      mFunction = new TF1J(name, dssFragRootWrapper,
                           mZLower, mZUpper, 6);
      // Set default parameters:
      mFunction->SetParameter(0, DSS); // This isn't actually used any more!
      mFunction->SetParameter(SCALE,  2.5);
      mFunction->SetParameter(PARTON, UQUARK);
      mFunction->SetParameter(ORDER,  mOrder);
      mFunction->SetParameter(HADRON, 1);
      mFunction->SetParameter(CHARGE, 1);
      mFunction->SetRange(mZLower, mZUpper);
   } // if
   return mFunction;
}

/*
 The c++ wrapper (DSSFragFunction) around the FORTRAN subroutine
 is itself called via a ROOT TF1 1D function.
 The TF1 takes 6 parameters (see DSSFragFunction for an explanation).
 */
DssFragmentation::DssFragmentation(int order, double meanPt2)
: mOrder(order)
, mMeanPt2(meanPt2)
, mFunction(NULL)
, mQBins(NULL) {
}

DssFragmentation::~DssFragmentation() {
   if(mFunction) {
      delete mFunction;
      mFunction = NULL;
   } // if
}

void DssFragmentation::SetFunctionParameters(int htype){
   double species(-1.e100);
   double charge(-1.e100);
   
   // Set the hadron type code
   switch(htype) {
      case 211://PIPLUS:
      case -211://PIMINUS:
      case 111://PIZERO:
         species = 1.;
         break;
      case 311://KZERO:
      case -311://KZEROBAR:
      case 321://KPLUS:
      case -321://KMINUS:
         species = 2.;
         break;
      case 2212://PROTON:
      case -2212://PBAR:
      case 2112://NEUTRON:
      case -2112://NBAR:
         species = 3.;
         break;
   } // switch
   
   // Set the hadron charge
   switch(htype) {
      case 211://PIPLUS:
      case 321://KPLUS:
      case 2212://PROTON:
         charge = 1.;
         break;
      case -211://PIMINUS:
      case -321://KMINUS:
      case -2212://PBAR:
         charge = -1.;
         break;
      case 111://PIZERO:
      case 311://KZERO:
      case -311://KZEROBAR:
      case 2112://NEUTRON:
      case -2112://NBAR:
         charge = 0.;
         break;
   } // switch
   
   Function()->SetParameter(HADRON, species); 
   Function()->SetParameter(CHARGE, charge);
}

void DssFragmentation::Insert(int hadron) {
   // Add the parton/Q2 table for this hadron
   std::vector<double> edges = mQBins->Edges();
   mTable.insert(std::make_pair(hadron,
                                KeyedLists(mQBins->Number(),
                                           edges.front(), edges.back())));   
   // Add the Q list for each parton in the new entry
   BOOST_FOREACH(int parton, Partons()) {
      HadronTable(hadron)->Insert(parton);
   } // for
}

bool DssFragmentation::Supports(const pythpart& parton) const {
   return std::find(Partons().begin(),
                    Partons().end(),
                    parton.type) not_eq Partons().end();
}

void DssFragmentation::PopulateKeyedLists() {
   // Save the parton type
   double saved_parton = Function()->GetParameter(PARTON);
   double saved_scale = Function()->GetParameter(SCALE);
   
   // Create a Q2/parton table for each hadron
   BOOST_FOREACH(int hadron, Hadrons()) {
      Insert(hadron);
   } // for
   
   /*
    For each (Q2 bin/hadron/parton), integrate the FF over
    the entire z range for which it is valid.
    Then normalise by the sum of integrals for all hadrons in
    that Q2 bin from that parton, so each entry is the probability 
    for that parton to fragment into that hadron type at that Q2.
    */
   BOOST_FOREACH(int parton, Partons()) { // Loop over parton types
      for(int qbin(0); qbin < mQBins->Number(); qbin++) { // Loop over Q2 bins
         ComputeHadronProbabilities(parton, qbin);
      } // for
   } // BOOST_FOREACH
   
   Function()->SetParameter(PARTON, saved_parton);
   Function()->SetParameter(SCALE,  saved_scale);
}

void DssFragmentation::ComputeHadronProbabilities(int parton, int qbin) {
   // All hadron tables have the same bin width, so just
   // pick one to get the bin size.
   double binLowEdge = mQBins->Edges().at(qbin);
   Function()->SetParameter(SCALE,  binLowEdge);
   Function()->SetParameter(PARTON, parton);
   
   double zLower(0.05), zUpper(1.);
   
   // DSS gives individual flavor/charge.
   // Calculate the integral of the FF for each hadron and accumulate
   // the total so we can normalise to 1 afterward.
   double sumIntegral(0.);
   BOOST_FOREACH(int hadron, Hadrons()) {
      // Set the function parameters for this Q2, parton and
      // compute the integral in this bin.
      SetFunctionParameters(hadron);
      // Change the epsilon used in integration from the default of 1.e-12
      // to 1.e-6. This saves significant time in intialisation. The z
      // distributions are all smooth, so the change in epsilon does not
      // seem to significantly impact the integral.
      double integral = Function()->Integral(zLower, zUpper,
                                             (double*)NULL, 1.e-6);
      // Set the value for this bin.
      // Accumulate the integral so we can normalise later.
      HadronTable(hadron)->Set(parton, qbin, integral);
      sumIntegral += integral;
   } // for
   
   // Find total for all hadrons from this parton at this Q2
   // and normalise to give probability.
   if(fabs(sumIntegral) < 1.e-5) {
      // Oops, don't want to divide by zero...
      std::cerr << "Hadron probability integral for (parton, Q) = ("
      << parton << ", " << qbin << ") == 0!" << std::endl;
      SetErrorValues(parton, qbin);
   } // else
   else {
      // Normalise:
      BOOST_FOREACH(int hadron, Hadrons()) {
         double probability =
         HadronTable(hadron)->At(parton, qbin) / sumIntegral;
         HadronTable(hadron)->Set(parton, qbin, probability);
      } // for
   }  // if
}

void DssFragmentation::SetErrorValues(int q, int qbin) {
   std::cerr <<
   " ERROR in DssFragmentation - sumIntegral = 0, selecting pion"
   << std::endl;
   mTable[211].Set(q, qbin, 1. / 3.);
   mTable[-211].Set(q, qbin, 1. / 3.);
   mTable[111].Set(q, qbin, 1. / 3.);
   mTable[321].Set(q, qbin, 0.);
   mTable[-321].Set(q, qbin, 0.);
   mTable[311].Set(q, qbin, 0.);
   mTable[-311].Set(q, qbin, 0.);
   mTable[2212].Set(q, qbin, 0.);
   mTable[-2212].Set(q, qbin, 0.);
   mTable[2112].Set(q, qbin, 0.);
   mTable[-2112].Set(q, qbin, 0.);
}

const std::vector<int>& DssFragmentation::Hadrons() const {
   static std::vector<int> hadrons;
   if(hadrons.empty()) {
      int array[11] = {
         211, -211, 111, 311, -311, 321, -321,
         2212, -2212, 2112, -2112
      };
      hadrons.assign(array, array + 11);
   } // if
   return hadrons;
}

const std::vector<int>& DssFragmentation::Partons() const {
   static std::vector<int> partons;
   if(partons.empty()) {
      int array[11] = {
         GLUON,UQUARK,UBAR,DQUARK,DBAR,SQUARK,SBAR,
         CQUARK,CBAR,BQUARK,BBAR
      };
      partons.assign(array, array + 11);
   } // if
   return partons;
}

KeyedLists* DssFragmentation::HadronTable(int hadron) {
   if(mTable.empty()) {
      PopulateKeyedLists();
   } // if
   Table::iterator i = mTable.find(hadron);
   if(i not_eq mTable.end()) {
      return &(i->second);
   } // if
   return NULL;
}

double DssFragmentation::Evaluate(const pythpart& parton,
                                  const Hadron& hadron,
                                  double scale) {
   Function()->SetParameter(SCALE,  ::sqrt(scale));
   Function()->SetParameter(PARTON, kkp::fromPdg(parton.type));
   
   SetFunctionParameters(hadron.type);
   return Function()->Eval(hadron.z);
}

int DssFragmentation::ThrowHadronType(const pythpart& parton,
                                      double Q2) {
   if(INVALID_HADRON not_eq mHadron) {
      return mHadron;
   } // if
   std::vector<int> hadrons = Hadrons();
   
   double base(0.);
   
   int chosenHadron(INVALID_HADRON);
   const double select = gRandom->Uniform(1.);
   BOOST_FOREACH(int hadron, hadrons) {
      double val = HadronTable(hadron)->Interpolate(kkp::fromPdg(parton.type), sqrt(Q2));
      base += val;
      if(base > select) {
         chosenHadron = hadron;
         break;
      } // if
   } // BOOST_FOREACH
     // JPSI is the last
   if(INVALID_HADRON == chosenHadron) {
      std::cerr << "Threw invalid hadron type... using pi0" << std::endl;
      chosenHadron = 111;//PIZERO;
   } // if
   
   return chosenHadron;
}

std::vector<Hadron> DssFragmentation::Fragment(const pythpart& parton,
                                               double Q2) {
   if(Q2 < 0.) {
      double pt2 = pow(parton.Pt(), 2.);
      if(pt2 < 2.) pt2 = 2.;
      Q2 = pt2;
   } // if
   
   int partontype = kkp::fromPdg(parton.Type());
   Function()->SetParameter(PARTON, partontype);
   Function()->SetParameter(SCALE,  ::sqrt(Q2));
   
   int h = ThrowHadronType(parton, Q2);
   Hadron hadron;
   hadron.type = h;
   Function()->SetRange(mZLower, mZUpper);
   hadron.z = Function()->GetRandom(mZLower, mZUpper);
   hadron.phigen = gRandom->Uniform(-TMath::Pi(), TMath::Pi());
   const double jt_sigma = sqrt(mMeanPt2/2.0);  
   double jt; 
   // Throw jt according to an exponential distribution such
   // that mean jt^2 is 0.2
   do {
      jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
   } while(jt>(3.0*jt_sigma));
   // Randomly generate phi.
   TVector3 mom(jt, 0., 0.);
   
   // TVector3::SetPhi rotates phi leaving magnitude and theta unchanged.
   mom.SetPhi(hadron.phigen);
   hadron.px = mom.Px();
   hadron.py = mom.Py();
   
   hadron.mParton = parton;
   hadron.mFrag = Evaluate(parton, hadron, Q2);
   std::vector<Hadron> products;
   products.push_back(hadron);
   return products;
}

bool DssFragmentation::Initialise() {
   mQBins.reset(new UniformBins(230, sqrt(MinimumScale()),
                                sqrt(MaximumScale())));
   return mQBins.get();
}

} // namespace tppmc
